#!/usr/bin/env python
#coding:utf-8

# 识别多边形
import rospy
import cv2
from geometry_msgs.msg import Twist
from sensor_msgs.msg import CompressedImage
from cv_bridge import CvBridge
import math
import numpy as np


class ShapeDetector:
    # 初始化类
    def __init__(self):
        # 初始化图形类型为不可识别
        self.shape = "unrecognized image"
        # 图形顶点集置空
        self.approx = []
        # 初始化该图形的周长为0
        self.peri = 0

    # 计算欧式距离(主要作用通过计算边长区分菱形和长方形)
    def distance(self, x1, y1, x2, y2):
        return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)

    def detect(self, c, frame):
        # cv2.arcLength函数返回周长
        self.peri = cv2.arcLength(c, True)

        # cv2.approxPolyDP用多边形取拟合，返回的是顶点的列表
        self.approx = cv2.approxPolyDP(c, 0.02 * self.peri, True)
        # self.approx = cv2.approxPolyDP(c, 0.2, True)
        d = len(self.approx)
        for i in range(d):
            frame2 = cv2.line(frame, (self.approx[i][0][0], self.approx[i][0][1]),
                        (self.approx[(i + 1) % d][0][0], self.approx[(i + 1) % d][0][1]), (255, 0, 0), 3)
        cv2.circle(frame2, (self.approx[i][0][0], self.approx[i][0][1]), 4, (0, 0, 255), -1)

        # 3个顶点，三角形
        if d == 3:
            self.shape = "triangle"
        # 同理，四个顶点，四边形
        elif d == 4:
            # 计算相邻两边的长度，做差判在误差范围内是否相等
            dist1 = self.distance(self.approx[0][0][0], self.approx[0][0][1], self.approx[1][0][0],
                            self.approx[1][0][1])
            dist2 = self.distance(self.approx[0][0][0], self.approx[0][0][1], self.approx[3][0][0],
                            self.approx[3][0][1])
            result = math.fabs(dist1 - dist2)
            # 误差小于10，可近似认为相等，为菱形
            if result <= 10:
                self.shape = "rhombus"
            else:
                self.shape = "rectangle"
        # 五边形
        elif d == 5:
            self.shape = "pentagon"
        # 六边形
        elif d == 6:
            self.shape = "hexagon"
        # 圆
        else:
            self.shape = "unrecognized"

        # 返回形状
        return self.shape, d, frame2



def main():
    # ROS节点初始化
    rospy.init_node('polygon_publisher', anonymous=True)

    # 创建两个Publisher，发布名为/vision/camera/image_raw的topic，消息类型为sensor_msgs::Image，队列长度10
    img_raw_pub = rospy.Publisher('/vision/camera/image_raw',CompressedImage, queue_size=10)
    img_hd_pub = rospy.Publisher('/vision/camera/image_handle',CompressedImage, queue_size=10)

    # 创建一个Publisher，发布名为/polygon/cmd_vel的topic，消息类型为geometry_msgs::Twist，队列长度10
    polygon_vel_pub = rospy.Publisher('/vision/cmd/shape', Twist, queue_size=10)

    bridge = CvBridge()

    cap = cv2.VideoCapture(2) 

    # 创建压缩图像
    msg = CompressedImage()

    rate = rospy.Rate(10)
    while not rospy.is_shutdown():

        # 获取摄像头图像
        ret, frame = cap.read()

        # 发布原始图像
        msg.header.stamp = rospy.Time.now()
        msg.format = "jpeg"
        msg.data = np.array(cv2.imencode('.jpg', frame)[1]).tostring()
        img_raw_pub.publish(msg)

        # 将图片转换为灰度图片
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # 高斯滤波,图像平滑处理
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)

        # 根据阈值，将灰度图片转化为黑白两色图片
        thresh = cv2.threshold(blurred, 250, 255, cv2.THRESH_BINARY)[1]

        # 返回图片和图中轮廓信息，查找最大轮廓返回到cnts中
        cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
        if len(cnts):
            area = []
            for i in range(len(cnts)):
                area.append(cv2.contourArea(cnts[i]))
            max_idx = np.argmax(np.array(area))

            # 创建一个识别器实例
            sd = ShapeDetector()
            # 对轮廓进行处理
            n_shape, point, frame2 = sd.detect(cnts[max_idx], frame)

            if point == 3 or point == 4 or point == 5:
                # 初始化geometry_msgs::Twist类型的消息
                vel_msg = Twist()
                vel_msg.linear.x = float(point)
                # 发布消息
                polygon_vel_pub.publish(vel_msg)
                rospy.loginfo("%0.2f", vel_msg.linear.x)

            else:
                # 初始化geometry_msgs::Twist类型的消息
                vel_msg = Twist()
                vel_msg.linear.x = float(0)
                # 发布消息
                polygon_vel_pub.publish(vel_msg)
                rospy.loginfo("%0.2f", vel_msg.linear.x)

        else:
            frame2 = frame

        # 发布处理后的图像
        msg.header.stamp = rospy.Time.now()
        msg.format = "jpeg"
        msg.data = np.array(cv2.imencode('.jpg', frame2)[1]).tostring()
        img_hd_pub.publish(msg)
        
        rate.sleep()



if __name__ == "__main__":
    try:
        main()
    except rospy.ROSInterruptException:
        pass


